from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import base64
import os
import json


class CryptoUtils:
    def __init__(self, key=None):
        """
        初始化加密工具
        :param key: 加密密钥，如果为None则生成新密钥
        """
        if key is None:
            # 生成新的密钥
            self.key = Fernet.generate_key()
        else:
            self.key = key
        self.cipher_suite = Fernet(self.key)

    def generate_key_from_password(self, password, salt=None):
        """
        从密码生成密钥
        :param password: 密码字符串
        :param salt: 盐值，如果为None则生成新盐值
        :return: 密钥
        """
        if salt is None:
            salt = os.urandom(16)

        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )
        key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
        return key, salt

    def encrypt_data(self, data):
        """
        加密数据
        :param data: 要加密的数据（字符串或字典）
        :return: 加密后的数据
        """
        if isinstance(data, dict):
            data = json.dumps(data, ensure_ascii=False)

        if isinstance(data, str):
            data = data.encode("utf-8")

        encrypted_data = self.cipher_suite.encrypt(data)
        return base64.urlsafe_b64encode(encrypted_data).decode("utf-8")

    def decrypt_data(self, encrypted_data):
        """
        解密数据
        :param encrypted_data: 加密的数据
        :return: 解密后的数据
        """
        if isinstance(encrypted_data, str):
            encrypted_data = encrypted_data.encode("utf-8")

        encrypted_data = base64.urlsafe_b64decode(encrypted_data)
        decrypted_data = self.cipher_suite.decrypt(encrypted_data)

        # 尝试解析为JSON，如果失败则返回字符串
        try:
            return json.loads(decrypted_data.decode("utf-8"))
        except json.JSONDecodeError:
            return decrypted_data.decode("utf-8")

    def get_key(self):
        """
        获取当前密钥
        :return: 密钥字符串
        """
        return base64.urlsafe_b64encode(self.key).decode("utf-8")

    def set_key(self, key):
        """
        设置密钥
        :param key: 密钥字符串
        """
        if isinstance(key, str):
            key = key.encode("utf-8")
        self.key = base64.urlsafe_b64decode(key)
        self.cipher_suite = Fernet(self.key)


# 使用示例
def demo_encryption():
    """演示加密解密功能"""
    print("=== 加密解密演示 ===")

    # 创建加密工具实例
    crypto = CryptoUtils()

    # 示例数据
    test_data = {"name": "张三", "age": 25, "skills": ["Python", "JavaScript", "SQL"]}

    print(f"原始数据: {test_data}")

    # 加密数据
    encrypted = crypto.encrypt_data(test_data)
    print(f"加密后: {encrypted}")

    # 解密数据
    decrypted = crypto.decrypt_data(encrypted)
    print(f"解密后: {decrypted}")

    # 保存密钥
    key = crypto.get_key()
    print(f"密钥: {key}")

    # 使用保存的密钥创建新的加密工具
    new_crypto = CryptoUtils()
    new_crypto.set_key(key)

    # 使用新工具解密
    decrypted_again = new_crypto.decrypt_data(encrypted)
    print(f"使用新工具解密: {decrypted_again}")


if __name__ == "__main__":
    demo_encryption()
